import os
import json
import argparse

import numpy as np

import DocTer.fuzzer.util as util
from DocTer.fuzzer.fuzzer import Fuzzer
from DocTer.fuzzer.fuzzer_config import FuzzerConfig
from utils.encoder import GreatEncoder


def generate_inputs(constraint_path, dtype_path, output_path, fuzz_optional_p=0.2, guided_mutate=True,
                    mutate_p=0.4, max_iter=10):
    target_config = util.read_yaml(os.path.abspath(constraint_path))
    configs = {
        "fuzz_optional_p": fuzz_optional_p,
        "guided_mutate": guided_mutate,
        "mutate_p": mutate_p,
        "max_iter": max_iter,
        'target_config': target_config,
        'dtype_config': util.read_yaml(os.path.abspath(dtype_path))
    }

    fuzzer_config = FuzzerConfig(configs)
    fuzzer = Fuzzer(fuzzer_config)
    all_name = target_config.get("title")

    full_output_path = os.path.join(output_path, all_name + "_seeds.json")
    test_cases = []
    # save numpy arrays to .npz file
    nd_array_names = {}

    with open(full_output_path, "w+") as f:
        cnt = 0
        while cnt < max_iter:
            new_seed = fuzzer._generate_new_input(True)
            param_list = new_seed.keys()
            # pre-handle new seed, especially ndarray
            for param in param_list:
                val = new_seed[param]
                if type(val) is np.ndarray:
                    array_name = param + "_" + str(cnt)
                    nd_array_names[array_name] = val
                    new_seed[param] = array_name
            test_cases.append(new_seed)
            cnt += 1
        # write new seeds into json
        json.dump(test_cases, f, cls=GreatEncoder)
        # write numpy arrays into npz
        np.savez(os.path.join(output_path, all_name + "_seeds"), **nd_array_names)


if __name__ == "__main__":
    current_path = os.path.dirname(os.path.abspath(__file__))
    parser = argparse.ArgumentParser()
    parser.add_argument("--dtype_path", "-d", default=os.path.join(current_path, "DocTer/pytorch_dtypes.yml"))
    parser.add_argument("--constraint_path", "-c", default=os.path.join(current_path, "DocTer/constraints/pytorch"
                                                                                      "/torch.nn.functional.conv2d.yaml"))
    parser.add_argument("--output_path", "-o", default="/home/ubuntu/Ascend/ascend_testing/model_transfer"
                                                       "/model_transfer_source/seed/torch.nn.functional.conv2d")
    parser.add_argument("--max_iter", type=int, default=1500)

    args = parser.parse_args()
    generate_inputs(args.constraint_path, args.dtype_path, args.output_path, max_iter=args.max_iter)

